fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics - #5956

source/tests/pt_expt/test_compress_min_nbor_dist.py (1)

66-70: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Add execution-path tests for forced recomputation.

This test only confirms parser state. Add focused tests that verify enable_compression ignores a saved value, assigns the value returned by UpdateSel.get_min_nbor_dist, and rejects a missing training_script.

As per coding guidelines, use pytest for single test cases instead of the full test suite.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@source/tests/pt_expt/test_compress_min_nbor_dist.py` around lines 66 - 70,
Extend the compression tests around enable_compression to cover forced
recomputation: verify the saved minimum-neighbor-distance value is ignored, the
result from UpdateSel.get_min_nbor_dist is assigned, and a missing
training_script is rejected. Use focused pytest test cases with mocked
dependencies rather than invoking the full test suite.

Source: Coding guidelines

deepmd/pt_expt/entrypoints/compress.py:43

• _read_saved_min_nbor_dist assumes model_dict["@variables"] is a dict by doing (model_dict.get("@variables") or {}).get(...). If @variables is present but not a mapping (corrupt/older metadata, or other producers), this will raise AttributeError during compress. Other code paths in the repo defensively guard @variables with isinstance(..., dict) (e.g. pt_expt DPA4 normalization).

        return float(min_nbor_dist), "the model file"
    min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
    if min_nbor_dist is not None:
        return float(min_nbor_dist), "the model file (@variables)"
    return None, ""

deepmd/main.py:677

• --recompute-min-nbor-dist is added to the top-level dp compress parser, so it will be accepted for all backends. However, only the PyTorch Exportable backend actually reads/uses recompute_min_nbor_dist; other backends ignore it (e.g. deepmd/pt/entrypoints/main.py and deepmd/jax/entrypoints/main.py don't pass it through). This can mislead users because the CLI will accept the flag but it will have no effect unless the backend is PyTorch Exportable.

Consider validating at dispatch time (or in each backend entrypoint) that this flag is only allowed with the PyTorch Exportable backend, and error out otherwise.

    parser_compress.add_argument(
        "--recompute-min-nbor-dist",
        action="store_true",
        help="(Supported backend: PyTorch Exportable) Ignore the minimal neighbor "
        "distance saved in the model and recompute it from the training data. "
        "Requires -t,--training-script",
    )

deepmd/pt_expt/entrypoints/compress.py:41

• model_dict.get("@variables") is assumed to be a dict, but if a malformed/older model file stores a truthy non-mapping (e.g. a list/string), (model_dict.get("@variables") or {}).get(...) will raise AttributeError. Since this data comes from disk, it’s safer to guard the type and treat non-dict values as “not present” (or raise a clearer ValueError).

    min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
    if min_nbor_dist is not None:
        return float(min_nbor_dist), "the model file (@variables)"
    return None, ""

deepmd/pt_expt/utils/neighbor_stat.py:82

• This change routes neighbor-stat computation through AutoBatchSize.execute_all, but there’s no unit-level regression test to ensure batching is actually invoked (and keeps working) for the pt_expt backend. Consider adding a lightweight test (e.g. monkeypatch AutoBatchSize.execute_all to assert it’s called) so future refactors don’t accidentally revert to whole-set execution.

                minrr2, max_nnei = self.auto_batch_size.execute_all(
                    self._execute,
                    data_set_data["coord"].shape[0],
                    data_set.get_natoms(),
                    data_set_data["coord"],